{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "episodes: 0 to 100, average_reward: 31.109, exploration: 0.598\n",
      "episodes: 100 to 200, average_reward: 109.020, exploration: 0.132\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-1-792381b04b12>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m    241\u001b[0m     \u001b[0menv\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgym\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmake\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'CartPole-v1'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    242\u001b[0m     \u001b[0magent\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mAgent\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0menv\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 243\u001b[0;31m     \u001b[0magent\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    244\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-1-792381b04b12>\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    223\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    224\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madd\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstate\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maction\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreward\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnext_state\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdone\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 225\u001b[0;31m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mqnet\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbatch_train\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbatch_size\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    226\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    227\u001b[0m                 \u001b[0;31m# update target network\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import gym\n",
    "\n",
    "class ExpReplay():\n",
    "    def __init__(self, e_max=15000, e_min=100):\n",
    "        self._max = e_max # maximum number of experiences\n",
    "        self._min = e_min # minimum number of experiences for training\n",
    "        self.exp = {'state':[], 'action':[], 'reward':[], 'next_state':[], 'done':[]} # total experiences the Agent stored\n",
    "        \n",
    "    def get_max(self):\n",
    "        \"\"\"return the maximum number of experiences\"\"\"\n",
    "        return self._max\n",
    "    \n",
    "    def get_min(self):\n",
    "        \"\"\"return the minimum number of experiences\"\"\"\n",
    "        return self._min\n",
    "    \n",
    "    def get_num(self):\n",
    "        \"\"\"return the curren number of experiences\"\"\"\n",
    "        return len(self.exp['state'])\n",
    "    \n",
    "    def get_batch(self, batch_size=64):\n",
    "        \"\"\"random choose a batch of experiences for training\"\"\"\n",
    "        idx = np.random.choice(self.get_num(), size=batch_size, replace=False)\n",
    "        state = np.array([self.exp['state'][i] for i in idx])\n",
    "        action = [self.exp['action'][i] for i in idx]\n",
    "        reward = [self.exp['reward'][i] for i in idx]\n",
    "        next_state = np.array([self.exp['next_state'][i] for i in idx])\n",
    "        done = [self.exp['done'][i] for i in idx]\n",
    "        return state, action, reward, next_state, done\n",
    "        \n",
    "    def add(self, state, action, reward, next_state, done):\n",
    "        \"\"\"add single experience\"\"\"\n",
    "        if self.get_num()>self.get_max():\n",
    "            del self.exp['state'][0]\n",
    "            del self.exp['action'][0]\n",
    "            del self.exp['reward'][0]\n",
    "            del self.exp['next_state'][0]\n",
    "            del self.exp['done'][0]\n",
    "        self.exp['state'].append(state)\n",
    "        self.exp['action'].append(action)\n",
    "        self.exp['reward'].append(reward)\n",
    "        self.exp['next_state'].append(next_state)\n",
    "        self.exp['done'].append(done)\n",
    "        \n",
    "        \n",
    "class TNET():\n",
    "    \"\"\"\n",
    "    Target network is for calculating the maximum estimated Q-value in given action a.\n",
    "    \"\"\"\n",
    "    def __init__(self, in_units, out_units, hidden_units=250):\n",
    "        self.in_units = in_units\n",
    "        self.out_units = out_units\n",
    "        self.hidden_units = hidden_units\n",
    "        self._model()\n",
    "        \n",
    "    def _model(self):\n",
    "        with tf.variable_scope('tnet'):\n",
    "            self.x = tf.placeholder(tf.float32, shape=(None, self.in_units))\n",
    "            \n",
    "            W1=tf.get_variable('W1', shape=(self.in_units, self.hidden_units), initializer=tf.random_normal_initializer())\n",
    "            W2=tf.get_variable('W2', shape=(self.hidden_units, self.hidden_units), initializer=tf.random_normal_initializer())\n",
    "            W3=tf.get_variable('W3', shape=(self.hidden_units, self.out_units), initializer=tf.random_normal_initializer())\n",
    "            \n",
    "            b1=tf.get_variable('b1', shape=(self.hidden_units), initializer=tf.zeros_initializer())\n",
    "            b2=tf.get_variable('b2', shape=(self.hidden_units), initializer=tf.zeros_initializer())\n",
    " \n",
    "            h1=tf.nn.tanh(tf.matmul(self.x, W1)+b1)\n",
    "            h2=tf.nn.tanh(tf.matmul(h1, W2)+b2)\n",
    "            self.q=tf.matmul(h2, W3)\n",
    "\n",
    "            self.params = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='tnet')\n",
    "            \n",
    "            \n",
    "class QNET():\n",
    "    def __init__(self, in_units, out_units, exp, hidden_units=250):\n",
    "        # Target Network\n",
    "        self.tnet = TNET(in_units, out_units)\n",
    "        \n",
    "        # Q network architecture\n",
    "        self.in_units = in_units\n",
    "        self.out_units = out_units\n",
    "        self.hidden_units = hidden_units\n",
    "        self._model()\n",
    "        self._batch_learning_model()\n",
    "        self._tnet_update()\n",
    "        \n",
    "        # experience replay\n",
    "        self.exp = exp \n",
    "        \n",
    "    def _model(self):\n",
    "        \"\"\" Q-network architecture \"\"\"\n",
    "        with tf.variable_scope('qnet'):\n",
    "            self.x = tf.placeholder(tf.float32, shape=(None, self.in_units))\n",
    "            \n",
    "            W1 = tf.get_variable('W1', shape=(self.in_units, self.hidden_units), initializer=tf.random_normal_initializer())\n",
    "            W2 = tf.get_variable('W2', shape=(self.hidden_units, self.hidden_units), initializer=tf.random_normal_initializer())\n",
    "            W3 = tf.get_variable('W3', shape=(self.hidden_units, self.out_units), initializer=tf.random_normal_initializer())\n",
    "            \n",
    "            b1 = tf.get_variable('b1', shape=(self.hidden_units), initializer=tf.zeros_initializer())\n",
    "            b2 = tf.get_variable('b2', shape=(self.hidden_units), initializer=tf.zeros_initializer())\n",
    " \n",
    "            h1 = tf.nn.tanh(tf.matmul(self.x, W1)+b1)\n",
    "            h2 = tf.nn.tanh(tf.matmul(h1, W2)+b2)\n",
    "            self.q = tf.matmul(h2, W3)\n",
    "            \n",
    "    def _batch_learning_model(self):\n",
    "        \"\"\"For batch learning\"\"\"\n",
    "        with tf.variable_scope('qnet'):\n",
    "            # TD-target\n",
    "            self.target = tf.placeholder(tf.float32, shape=(None, ))\n",
    "            # Action index\n",
    "            self.selected_idx = tf.placeholder(tf.int32, shape=(None, 2))\n",
    "            # Q-value\n",
    "            self.selected_q = tf.gather_nd(self.q, self.selected_idx)\n",
    "            \n",
    "            self.params = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='qnet')\n",
    "            \n",
    "            # Q-network optimization alogrithms\n",
    "            loss = tf.losses.mean_squared_error(self.target, self.selected_q)\n",
    "            gradients = tf.gradients(loss, self.params)\n",
    "            self.train_opt = tf.train.AdamOptimizer(3e-4).apply_gradients(zip(gradients, self.params))\n",
    "\n",
    "    def _tnet_update(self):\n",
    "        \"\"\" Update Target network by using the parameters of Q-Network\"\"\"\n",
    "        with tf.variable_scope('qnet'):                        \n",
    "            self.update_opt = [t.assign(q) for t, q in zip(self.tnet.params, self.params)]\n",
    "    \n",
    "    def batch_train(self, batch_size=64):\n",
    "        \"\"\"Implement Double DQN Algorithm, batch training\"\"\"\n",
    "        if self.exp.get_num() < self.exp.get_min():\n",
    "            #The number of experiences is not enough for batch training\n",
    "            return\n",
    "\n",
    "        # get a batch of experiences\n",
    "        state, action, reward, next_state, done = self.exp.get_batch(batch_size)\n",
    "        state = state.reshape(batch_size, self.in_units)\n",
    "        next_state = next_state.reshape(batch_size, self.in_units)\n",
    "        \n",
    "        # get actions by Q-network\n",
    "        qnet_q_values = self.session.run(self.q, feed_dict={self.x:next_state})\n",
    "        qnet_actions = np.argmax(qnet_q_values, axis=1)\n",
    "        \n",
    "        # calculate estimated Q-values with qnet_actions by using Target-network\n",
    "        tnet_q_values = self.session.run(self.tnet.q, feed_dict={self.tnet.x:next_state})\n",
    "        tnet_q = [np.take(tnet_q_values[i], qnet_actions[i]) for i in range(batch_size)]\n",
    "        \n",
    "        # Update Q-values of Q-network\n",
    "        qnet_update_q = [r+0.95*q if not d else r for r, q, d in zip(reward, tnet_q, done)]\n",
    "        \n",
    "        # optimization\n",
    "        indices=[[i,action[i]] for i in range(batch_size)]\n",
    "        feed_dict={self.x:state, self.target:qnet_update_q, self.selected_idx:indices}\n",
    "        self.session.run(self.train_opt, feed_dict)\n",
    "        \n",
    "    def update(self):\n",
    "        \"\"\" for updatte target network\"\"\"\n",
    "        self.session.run(self.update_opt)\n",
    "        \n",
    "    def set_session(self, sess):\n",
    "        self.session = sess\n",
    "        \n",
    "    def get_action(self, state, e_rate):\n",
    "        \"\"\" for training stage of the Agent, exploitation or exploration\"\"\"\n",
    "        if np.random.random()<e_rate:\n",
    "            return np.random.choice(self.out_units)\n",
    "        else:\n",
    "            return np.argmax(self.session.run(self.q, feed_dict={self.x: state}))\n",
    "        \n",
    "        \n",
    "class Agent():\n",
    "    def __init__(self, env):\n",
    "        # set hyper parameters\n",
    "        self.max_episodes = 10000\n",
    "        self.max_actions = 10000\n",
    "        self.exploration_rate = 1.0\n",
    "        self.exploration_decay = 0.0001  \n",
    "        \n",
    "        # set environment\n",
    "        self.env = env\n",
    "        self.states = env.observation_space.shape[0]\n",
    "        self.actions = env.action_space.n\n",
    "        \n",
    "        # Experience Replay for batch learning\n",
    "        self.exp = ExpReplay()\n",
    "        # the number of experience per batch for batch learning\n",
    "        self.batch_size = 64 \n",
    "        \n",
    "        # Deep Q Network\n",
    "        self.qnet = QNET(self.states, self.actions, self.exp)\n",
    "        # For execute Deep Q Network\n",
    "        session = tf.InteractiveSession()\n",
    "        session.run(tf.global_variables_initializer())\n",
    "        self.qnet.set_session(session)\n",
    "        \n",
    "    def train(self):\n",
    "        # set hyper parameters\n",
    "        max_episodes = self.max_episodes\n",
    "        max_actions = self.max_actions\n",
    "        exploration_rate = self.exploration_rate\n",
    "        exploration_decay = self.exploration_decay\n",
    "        batch_size = self.batch_size\n",
    "        \n",
    "        # start training\n",
    "        record_rewards = []\n",
    "        for i in range(max_episodes):\n",
    "            total_rewards = 0\n",
    "            state = self.env.reset()\n",
    "            state = state.reshape(1, self.states)\n",
    "            for j in range(max_actions):\n",
    "                #self.env.render() # Uncomment this line to render the environment\n",
    "                action = self.qnet.get_action(state, exploration_rate)\n",
    "                next_state, reward, done, info = self.env.step(action)\n",
    "                next_state = next_state.reshape(1, self.states)\n",
    "                total_rewards += reward\n",
    "                \n",
    "                if done:\n",
    "                    self.exp.add(state, action, (reward-100), next_state, done)\n",
    "                    self.qnet.batch_train(batch_size)\n",
    "                    break\n",
    "                    \n",
    "                self.exp.add(state, action, reward, next_state, done)\n",
    "                self.qnet.batch_train(batch_size)\n",
    "                \n",
    "                # update target network\n",
    "                if (j%25)== 0 and j>0:\n",
    "                    self.qnet.update()\n",
    "                # next episode\n",
    "                state = next_state\n",
    "                \n",
    "            record_rewards.append(total_rewards)\n",
    "            exploration_rate = 0.001 + (exploration_rate-0.001)*np.exp(-exploration_decay*(i+1))\n",
    "            if i%100==0 and i>0:\n",
    "                average_rewards = np.mean(np.array(record_rewards))\n",
    "                record_rewards = []\n",
    "                print(\"episodes: %i to %i, average_reward: %.3f, exploration: %.3f\" %(i-100, i, average_rewards, exploration_rate))\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    env = gym.make('CartPole-v1')\n",
    "    agent = Agent(env)\n",
    "    agent.train()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
